/*
 * Copyright (c) 2010 Toni Spets <toni.spets@iki.fi>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <windows.h>
#include <psapi.h>
#include <stdio.h>

#include <stdint.h>
#include "config.h"

#include "inject.c"

#define CNC95_106br2    2306560
#define RA95_303_EXE      73728 /* launcher, real */
#define RA95_303_EXE_UDP  90112 /* launcher, RA303UDP version */
#define RA95_303_DAT    2375680 /* game */
#define SUN_203_EXE       81920 /* launcher */
#define SUN_203_GAME    3457296 /* launcher */
#define RA2_1006_EXE      57344 /* launcher */
#define RA2_1006_GAME   4485392 /* game */
#define RA2MD_1006_EXE  1314623 /* launcher */
#define RA2MD_1006_GAME 4813072 /* game */

DWORD WINAPI _IPX_Initialise();
DWORD WINAPI _IPX_Open_Socket95(DWORD);
DWORD WINAPI _IPX_Start_Listening95();
DWORD WINAPI _IPX_Get_Outstanding_Buffer95( void * );
DWORD WINAPI _IPX_Broadcast_Packet95( void *, DWORD );
DWORD WINAPI _IPX_Send_Packet95( void *, void *, DWORD, void *, void * );
DWORD WINAPI _IPX_Get_Connection_Number95();
DWORD WINAPI _IPX_Get_Local_Target95( DWORD *, DWORD *, DWORD *, DWORD * );
DWORD WINAPI _IPX_Close_Socket95( DWORD );
DWORD WINAPI _IPX_Shut_Down95();

WINAPI SOCKET WSAAPI fake_socket( int, int, int );
WINAPI int fake_bind( SOCKET, const struct sockaddr *, int );
WINAPI int fake_recvfrom( SOCKET, char *, int, int, struct sockaddr *, int * );
WINAPI int fake_sendto( SOCKET, const char *, int, int, const struct sockaddr *, int );
WINAPI int fake_getsockopt( SOCKET, int, int, char *, int * );
WINAPI int fake_setsockopt( SOCKET, int, int, const char *, int );

/* original func_replace_* functions courtesy of Steinar H. Gunderson <sgunderson@bigfoot.com>, rewritten by me */
void func_replace_offset(HANDLE hProcess, DWORD address, void* new_address)
{
    SIZE_T tmp = 0;
    if(!WriteProcessMemory(hProcess, (void*)address, &new_address, sizeof(DWORD), &tmp))
    {
        MessageBoxA(NULL, "Error: Memory patch failed!", "CnCNet", MB_OK);
        exit(1);
    }
}

void func_replace_hotpatch(HANDLE hProcess, DWORD address, void* new_address)
{
    unsigned char jmp_opcode = 0xe9;
    DWORD relative_addr = (DWORD)new_address - (DWORD)address - 5;
    SIZE_T tmp = 0;

    WriteProcessMemory(hProcess, (void*)address, &jmp_opcode, 1, &tmp);
    if(!WriteProcessMemory(hProcess, (void*)address + 1, &relative_addr, sizeof(DWORD), &tmp))
    {
        MessageBoxA(NULL, "Error: Memory patch failed!\n", "CnCNet", MB_OK);
        exit(1);
    }
}

void func_replace_hotpatch_call(HANDLE hProcess, DWORD address, void* new_address)
{
    SIZE_T tmp = 0;
    unsigned char jmp_opcode = 0xE8;
    unsigned char nop_opcode = 0x90;
    DWORD relative_addr = (DWORD)new_address - (DWORD)address - 5;

    WriteProcessMemory(hProcess, (void*)address, &jmp_opcode, 1, &tmp);
    WriteProcessMemory(hProcess, (void*)address + 1, &relative_addr, sizeof(DWORD), &tmp);
    WriteProcessMemory(hProcess, (void*)address + 1 + sizeof(DWORD), &nop_opcode, 1, &tmp);
}

char dllPath[MAX_PATH] = { 0 };

/* for injection trough launchers */
BOOL WINAPI fake_CreateProcessA(
  LPCTSTR lpApplicationName,
  LPTSTR lpCommandLine,
  LPSECURITY_ATTRIBUTES lpProcessAttributes,
  LPSECURITY_ATTRIBUTES lpThreadAttributes,
  BOOL bInheritHandles,
  DWORD dwCreationFlags,
  LPVOID lpEnvironment,
  LPCTSTR lpCurrentDirectory,
  LPSTARTUPINFO lpStartupInfo,
  LPPROCESS_INFORMATION lpProcessInformation
)
{
    /* workaround RA303UDP package launcher bat */
    if( _stricmp(lpCommandLine, "ra95.bat") == 0 )
    {
        lpCommandLine = "ra95.dat";
    }

    return CreateInjectedProcessA(lpApplicationName, lpCommandLine, lpProcessAttributes, lpThreadAttributes, bInheritHandles, dwCreationFlags|CREATE_SUSPENDED, lpEnvironment, lpCurrentDirectory, lpStartupInfo, lpProcessInformation, dllPath);
}

BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
{
    HANDLE hFile;
    HANDLE hProcess;
    char lpImagePath[MAX_PATH] = { 0 };
    char lpImageFileName[MAX_PATH] = { 0 };
    int size = 0;

    if(fdwReason == DLL_PROCESS_ATTACH)
    {
        config_init();

        hProcess = OpenProcess(PROCESS_SET_INFORMATION|PROCESS_QUERY_INFORMATION|PROCESS_VM_READ|PROCESS_VM_OPERATION|PROCESS_VM_WRITE, FALSE, GetCurrentProcessId());

        GetModuleFileName(NULL, lpImagePath, MAX_PATH);
        GetModuleBaseName(hProcess, NULL, lpImageFileName, MAX_PATH);

        GetModuleFileName(hinstDLL, dllPath, MAX_PATH);

        if( (hFile = CreateFile(lpImagePath, GENERIC_READ, FILE_SHARE_DELETE|FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)) )
        {
            size = GetFileSize(hFile, NULL);
            CloseHandle(hFile);

            if(_stricmp(lpImageFileName, "C&C95.exe") == 0)
            {
                if(size == CNC95_106br2)
                {
                    func_replace_offset(hProcess, 0x005B06EC, _IPX_Broadcast_Packet95);
                    func_replace_offset(hProcess, 0x005B06F0, _IPX_Close_Socket95);
                    func_replace_offset(hProcess, 0x005B06F4, _IPX_Get_Connection_Number95);
                    func_replace_offset(hProcess, 0x005B06F8, _IPX_Get_Local_Target95);
                    func_replace_offset(hProcess, 0x005B06FC, _IPX_Get_Outstanding_Buffer95);
                    func_replace_offset(hProcess, 0x005B0700, _IPX_Initialise);
                    func_replace_offset(hProcess, 0x005B0704, _IPX_Open_Socket95);
                    func_replace_offset(hProcess, 0x005B0708, _IPX_Send_Packet95);
                    func_replace_offset(hProcess, 0x005B070C, _IPX_Shut_Down95);
                    func_replace_offset(hProcess, 0x005B0710, _IPX_Start_Listening95);

                    SetProcessAffinityMask(hProcess, 0x01);
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports version 1.06b r2 of Command & Conquer, please upgrade", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else if(_stricmp(lpImageFileName, "RA95.exe") == 0 || _stricmp(lpImageFileName, "RA95.dat") == 0)
            {
                if(size == RA95_303_EXE || size == RA95_303_EXE_UDP)
                {
                    func_replace_hotpatch_call(hProcess, 0x00403D15, fake_CreateProcessA);
                }
                else if(size == RA95_303_DAT)
                {
                    func_replace_hotpatch(hProcess, 0x005E59E8, fake_bind);       /*  2 */
                    func_replace_hotpatch(hProcess, 0x005E5A00, fake_setsockopt); /*  7 */
                    func_replace_hotpatch(hProcess, 0x005E59D6, fake_recvfrom);   /* 17 */
                    func_replace_hotpatch(hProcess, 0x005E59D0, fake_sendto);     /* 20 */
                    func_replace_hotpatch(hProcess, 0x005E59FA, fake_setsockopt); /* 21 */
                    func_replace_hotpatch(hProcess, 0x005E59F4, fake_socket);     /* 23 */

                    SetProcessAffinityMask(hProcess, 0x01);
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports version 3.03 of Red Alert, please upgrade", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else if(_stricmp(lpImageFileName, "SUN.exe") == 0)
            {
                if(size == SUN_203_EXE)
                {
                    func_replace_hotpatch_call(hProcess, 0x00403F05, fake_CreateProcessA);
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports version 2.03 of Tiberian Sun, please upgrade", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else if(_stricmp(lpImageFileName, "GAME.exe") == 0)
            {
                if(size == SUN_203_GAME)
                {
                    func_replace_hotpatch(hProcess, 0x006B4D54, fake_bind);       /*  2 */
                    func_replace_hotpatch(hProcess, 0x006B4D48, fake_setsockopt); /*  7 */
                    func_replace_hotpatch(hProcess, 0x006B4D66, fake_recvfrom);   /* 17 */
                    func_replace_hotpatch(hProcess, 0x006B4D6C, fake_sendto);     /* 20 */
                    func_replace_hotpatch(hProcess, 0x006B4D60, fake_setsockopt); /* 21 */
                    func_replace_hotpatch(hProcess, 0x006B4D5A, fake_socket);     /* 23 */
                }
                else if(size == RA2_1006_GAME)
                {
                    func_replace_hotpatch(hProcess, 0x00781B9C, fake_bind);       /*  2 */
                    func_replace_hotpatch(hProcess, 0x00781B90, fake_setsockopt); /*  7 */
                    func_replace_hotpatch(hProcess, 0x00781BAE, fake_recvfrom);   /* 17 */
                    func_replace_hotpatch(hProcess, 0x00781BB4, fake_sendto);     /* 20 */
                    func_replace_hotpatch(hProcess, 0x00781BA8, fake_setsockopt); /* 21 */
                    func_replace_hotpatch(hProcess, 0x00781BA2, fake_socket);     /* 23 */
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports unencrypted (TFD 1.03 patch or free release) version 2.03 of Tiberian Sun and 1.006 of Red Alert 2.", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else if(_stricmp(lpImageFileName, "GAMEMD.exe") == 0)
            {
                if(size == RA2MD_1006_GAME)
                {
                    func_replace_hotpatch(hProcess, 0x007C8998, fake_bind);       /*  2 */
                    func_replace_hotpatch(hProcess, 0x007C898C, fake_setsockopt); /*  7 */
                    func_replace_hotpatch(hProcess, 0x007C89AA, fake_recvfrom);   /* 17 */
                    func_replace_hotpatch(hProcess, 0x007C89B0, fake_sendto);     /* 20 */
                    func_replace_hotpatch(hProcess, 0x007C89A4, fake_setsockopt); /* 21 */
                    func_replace_hotpatch(hProcess, 0x007C899E, fake_socket);     /* 23 */
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports version 1.006 of Red Alert 2: Yuri's Revenge, please upgrade", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else if(_stricmp(lpImageFileName, "RA2.exe") == 0)
            {
                if(size == RA2_1006_EXE)
                {
                    func_replace_hotpatch_call(hProcess, 0x00401B3E, fake_CreateProcessA);
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports version 1.006 of Red Alert 2, please upgrade", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else if(_stricmp(lpImageFileName, "RA2MD.exe") == 0)
            {
                if(size == RA2MD_1006_EXE)
                {
                    MessageBox(NULL, "Error: Support for Red Alert 2: Yuri's Revenge is still incomplete.", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
                else
                {
                    MessageBox(NULL, "Error: CnCNet only supports unencrypted (TFD 1.03 patch) version 1.006 of Red Alert 2: Yuri's Revenge, please upgrade", "CnCNet", MB_OK|MB_ICONERROR);
                    exit(1);
                }
            }
            else
            {
                char tmp[MAX_PATH];
                snprintf(tmp, MAX_PATH, "Error: CnCNet was loaded with an unknown game version\nLoader: %s (%d bytes)", lpImageFileName, size);
                MessageBox(NULL, tmp, "CnCNet", MB_OK|MB_ICONERROR);
                exit(1);
            }
        }
        else
        {
            MessageBox(NULL, "Error: CnCNet was unable to check the game version", "CnCNet", MB_OK|MB_ICONERROR);
            exit(1);
        }

    }

    return TRUE;
}
